package com.sikaryofficial.ai.service;

/**
 * @author : wuweihong
 * @desc : TODO  请填写你的功能描述
 * @date : 2025-11-04
 */


import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;

@Service
public class RagService {

	private final ChatModel chatModel;
	private final ChatClient chatClient;
	private final VectorStore vectorStore;
	private final DocumentService documentService;


	public RagService(ChatModel chatModel, VectorStore vectorStore, DocumentService documentService, ChatClient chatClient) {
		this.chatClient = chatClient;
		this.chatModel = chatModel;
		this.vectorStore = vectorStore;
		this.documentService = documentService;
	}



	@Value("classpath:/prompts/rag-prompt-template.st")
	private Resource ragPromptTemplate;

	private String conversationId = "wuweihong";

	public String processQuery(String query) {
		// 1. 检索相关文档
		List<Document> similarDocuments = vectorStore.similaritySearch(query);
		System.out.println("检索到相关文档：" + similarDocuments.size());

		// 2. 构建上下文
		String context = similarDocuments.stream()
				.map(Document::getFormattedContent)
				.collect(Collectors.joining("\n\n"));

		// 3. 构建提示词
		PromptTemplate promptTemplate = new PromptTemplate(ragPromptTemplate);
		Prompt prompt = promptTemplate.create(Map.of(
				"current_date", new Date().toLocaleString(),
				"input", query,
				"context", context
		));
		// 4. 调用LLM生成回答
		return chatClient.prompt(prompt)
				.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
				.call().content();
	}


	public SseEmitter streamProcessQuery(String query, String conversationId, SseEmitter emitter) {
		// 1. 检索相关文档
		List<Document> similarDocuments = vectorStore.similaritySearch(query);
		System.out.println("检索到相关文档：" + similarDocuments.size());

		// 2. 构建上下文
		String context = similarDocuments.stream()
				.map(Document::getFormattedContent)
				.collect(Collectors.joining("\n\n"));

		// 3. 构建提示词
		PromptTemplate promptTemplate = new PromptTemplate(ragPromptTemplate);
		Prompt prompt = promptTemplate.create(Map.of(
				"current_date", new Date().toLocaleString(),
				"input", query,
				"context", context
		));
		// 4. 调用LLM生成回答
		chatClient.prompt(prompt).advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId)).stream().chatResponse().subscribe(chatResponse -> {
			try {
				String content = chatResponse.getResult().getOutput().getText();
				// 发送sse事件
				emitter.send(org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event().data(content).id(Long.toString(System.currentTimeMillis())).build());
			} catch (IOException e) {
				emitter.completeWithError(e);
			}
		}, emitter::completeWithError, emitter::complete);
		emitter.onCompletion(() -> {
			try {
				emitter.send(org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event().id(Long.toString(System.currentTimeMillis())).build());
			} catch (IOException e) {
				emitter.completeWithError(e);
			}
		});
		emitter.onTimeout(() -> emitter.completeWithError(new IOException("SSE timeout")));
		return emitter;
	}

	@GetMapping(produces = MediaType.TEXT_EVENT_STREAM_VALUE)
	public Flux<String> stream(@RequestParam("question") String question) {
		return Flux.defer(() ->
						chatClient
								.prompt(question)
								// 引导模型使用工具；禁止输出思维过程或 <think> 标签
								// 显式传入从 MCP 发现的工具集合
								.stream()
								.content()
				)
				// 把源执行切换到弹性线程池，避免在 Netty 事件循环里阻塞
				.subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic())
				// 去掉空白 token 与 <think> 片段
				.filter(token -> token != null && !token.isBlank())
				.filter(token -> !token.contains("<think"))
				.filter(token -> !token.contains("</think>"))
				// 轻量聚合，让输出更连贯
				.bufferTimeout(25, java.time.Duration.ofMillis(120))
				.map(list -> String.join("", list));
	}

	public String getConversationId() {
		return UUID.randomUUID().toString();
	}

	public void initKnowledgeBase(String filePath) throws IOException {
		documentService.loadAndStoreDocuments(filePath);
	}
}
